#include "imagematch.h"
#include "sift/xform.h"
#include "sift/utils.h"
#include "imagefeature.h"
#include <boost/foreach.hpp>
#include "inputimage.h"
#include "stitcher.h"
#include "imagefeature.h"
#include <boost/lexical_cast.hpp>
#include "commandlineoptions.h"

using namespace sift;
using namespace cshotel;
using namespace std;
using namespace boost;
using namespace cv;

typedef pair<weak_ptr<ImageFeature>, weak_ptr<ImageFeature> > FeatureMatchPair;

ImageMatch::ImageMatch(boost::weak_ptr<InputImage> first, boost::weak_ptr<InputImage> second) : firstImage(first), secondImage(second), transformationMatrix(0) {
    if(shared_ptr<InputImage> image1 = firstImage.lock()) {
    if(shared_ptr<InputImage> image2 = secondImage.lock()) {
        foreach(shared_ptr<ImageFeature> feature, image1->getImageFeatures()) {
            foreach(weak_ptr<ImageFeature> match, feature->getFeatureMatches()) {
                if(shared_ptr<ImageFeature> matchingFeature = match.lock()) {
                    if(InputImage::getImageId(matchingFeature->getImage()) == image2->getId()) {
                        featureMatches[weak_ptr<ImageFeature>(feature)] = match;
                        break;
                    }
                }
            }
        }
    }
    }
}

ImageMatch::~ImageMatch() {
    if(transformationMatrix != 0) {
        cvReleaseMat(&transformationMatrix);
    }
}

CvMat* ImageMatch::getTransformationMatrix() {
    return transformationMatrix;
}

vector<weak_ptr<ImageFeature> >& ImageMatch::getInliers() {
    return inliers;
}

vector<weak_ptr<ImageFeature> >& ImageMatch::getOutliers() {
    return outliers;
}

map<weak_ptr<ImageFeature>, weak_ptr<ImageFeature> >& ImageMatch::getFeatureMatches() {
    return featureMatches;
}

weak_ptr<InputImage> ImageMatch::getFirstImage() {
    return firstImage;
}

weak_ptr<InputImage> ImageMatch::getSecondImage() {
    return secondImage;
}

void ImageMatch::doRANSAC() {
    int numMatchingFeatures = featureMatches.size();
    feature *features = new feature[numMatchingFeatures * 2];
    feature **featureInliers = 0;
    int counter = 0;
    
    // set up the features in the format required by the library
    foreach(FeatureMatchPair p, featureMatches) {
        if(shared_ptr<ImageFeature> firstFeature = (p.first.lock())) {
        if(shared_ptr<ImageFeature> secondFeature = (p.second.lock())) {
            features[counter] = firstFeature->getSIFTFeature();
            features[counter].feature_data = new weak_ptr<ImageFeature>(p.first);
            features[counter + numMatchingFeatures] = secondFeature->getSIFTFeature();
            features[counter].fwd_match = &(features[counter + numMatchingFeatures]);
            ++counter;
        }
        }
    }
    if(counter != numMatchingFeatures) {
        dout("ERROR, counter != numMatchingFeatures" << endl);
    }
    int numFeatureInliers = 0;
    
    // calculate the homography using RANSAC
    transformationMatrix = ransac_xform(features, numMatchingFeatures, FEATURE_FWD_MATCH, lsq_homog, 4, 4000, homog_xfer_err, 3.0, &featureInliers, &numFeatureInliers);
    
    // store the inliers and outliers
    for(int ii = 0; ii < numFeatureInliers; ++ii) {
        inliers.push_back(*(static_cast<weak_ptr<ImageFeature>* >(featureInliers[ii]->feature_data)));
    }
    
    foreach(FeatureMatchPair p, featureMatches) {
        bool found = false;
        foreach(weak_ptr<ImageFeature> f, inliers) {
            if(!(f < p.first) && !(p.first < f)) {
                found = true;
                break;
            }
        }
        if(!found) {
            outliers.push_back(p.first);
        }
    }
    
    for(int ii = 0; ii < numMatchingFeatures; ++ii) {
        delete (static_cast<weak_ptr<ImageFeature>* >(features[ii].feature_data));
    }
    delete[] features;
    free(featureInliers);
}

// see the autopano paper by Matthew Brown and David Lowe
bool ImageMatch::doProbabilisticCheck() {
    return inliers.size() > 0.2 * featureMatches.size();
}

// write images visualising the match
void ImageMatch::writeDebugImages() {
    typedef pair<weak_ptr<ImageFeature>, weak_ptr<ImageFeature> > FeaturePair;
    vector<FeaturePair> features;
    
    // generate inlier image
    foreach(weak_ptr<ImageFeature> feature, inliers) {
        features.push_back(FeaturePair(feature, featureMatches[feature]));
    }
    InputImage::drawFeatureMatchImage(firstImage, secondImage, features, "inliers_");
    
    // generate outlier image
    features.clear();
    foreach(weak_ptr<ImageFeature> feature, outliers) {
        features.push_back(FeaturePair(feature, featureMatches[feature]));
    }
    InputImage::drawFeatureMatchImage(firstImage, secondImage, features, "outliers_");
}
